#!/usr/bin/env python

# This should work as an TEMPLATE for all future moveit related python code

import sys
import rospy
from yaml import load
try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader

import moveit_commander
import moveit_msgs.msg
import geometry_msgs.msg
import aubo_i5_hand_eye_cali.msg as cali_msg
import aubo_i5_hand_eye_cali.srv as cali_srv
import tf
import numpy
from threading import Thread


# This node accepts position commands of the manipulator with respect to
# the predefined table frame, and convert the command into base_link frame
# and send it to the robot to execute.

# note currently only pose vertical to the table is supported.

class MoveInTableNode(Thread):  #MoveInTableNode is a thread.

    def __init__(self):

        Thread.__init__(self)

        # read config files
        with open("../config/"+"aubo.yaml") as conf_file:
            config = load(conf_file, Loader=Loader)
            conf = config['hand_eye_conf']

            self.move_group_name = conf['move_group_name']           
            print "move_group_name: %s" % conf['move_group_name']
            self.cmd_ref_frame = conf['cmd_ref_frame']
            print "cmd_ref_frame: %s" % conf['cmd_ref_frame']
            self.world_frame = conf['world_frame']
            print "world_frame: %s" % conf['world_frame']            
            self.move_group_end_effector = conf['end_effector_name']
            print "move_group_end_effector: %s" % conf['end_effector_name']                  
            self.move_group_board_frame = conf['cali_board_name']
            print "move_group_board_frame: %s" % conf['cali_board_name']                  

        self.robot = moveit_commander.RobotCommander()
        self.scene = moveit_commander.PlanningSceneInterface()

        # the moveit group name defined according to the moveit config
        self.group = moveit_commander.MoveGroupCommander(self.move_group_name)
        self.group.set_planner_id("RRTConnectkConfigDefault")

        # create a DisplayTrajectory publisher for RVIZ to visualize
        self.display_trajectory_publisher = rospy.Publisher('/move_group/display_planned_path', moveit_msgs.msg.DisplayTrajectory)

        # sending planning command to ee_link
        self.group.set_end_effector_link(self.move_group_end_effector)

        self.frame_listener =  tf.TransformListener()


        # link6 to tool0 transformation
        self.tool0_to_ee_trans = []
        self.tool0_to_ee_rot = []

        pass

    def run(self):

        # listen to table frame broadcaster
        try:
            # table frame with respect to base_link

            rospy.loginfo('Listening to table frame tf...')

            # link_6 with respect to tool0
            self.frame_listener.waitForTransform(self.move_group_board_frame, self.move_group_end_effector ,rospy.Time(), rospy.Duration(4.0))
            (link6_trans, link6_rot)  = self.frame_listener.lookupTransform(self.move_group_board_frame , self.move_group_end_effector , rospy.Time(0))

            self.tool0_to_link6_trans = link6_trans
            self.tool0_to_link6_rot = link6_rot
#            print 'tool0_to_link6_trans:'
#            print link6_trans
#            print 'tool0_to_link6_rot:'
#            print link6_rot

            self.frame_listener.waitForTransform(self.world_frame, self.cmd_ref_frame,rospy.Time(), rospy.Duration(4.0))
            (world_trans, world_rot)  = self.frame_listener.lookupTransform(self.world_frame, self.cmd_ref_frame, rospy.Time(0))

            self.world_trans = world_trans
            self.world_rot = world_rot

            rospy.Service("move_arm_service",cali_srv.MoveArmRPYSrv, self.onMoveArmService)

            print "Waiting for command msgs.. "

            rospy.spin()

        except KeyboardInterrupt:
            pass
        except (tf.LookupException, tf.ConnectivityException):
            pass


        pass

    # input: a PickPoint msg data
    # output: the target Pose msg
    def pickPointMsgToTargetPose(self, msg):
        ## euler angle Z_Y_X
        target_rot = tf.transformations.quaternion_multiply(
                             tf.transformations.quaternion_about_axis(msg.theta, (0,0,1)),
                             tf.transformations.quaternion_about_axis(msg.beta, (0,1,0))
                             )
        target_rot = tf.transformations.quaternion_multiply(
                        target_rot,
                        tf.transformations.quaternion_about_axis(msg.alpha, (1,0,0))
                        )

        print 'target_rot_euler',tf.transformations.euler_from_quaternion(target_rot)

        target_trans = [msg.x, msg.y, msg.z]

        rospy.loginfo("======= pickPointMsgToTargetPose ======")
        rospy.loginfo("Converting incoming PickPoint command: %s %f", target_trans.__str__(), msg.theta)
#        print 'target_in_table: \n', target_trans

        self.frame_listener.waitForTransform(self.world_frame, self.cmd_ref_frame,rospy.Time(), rospy.Duration(1.0))
        (world_trans, world_rot)  = self.frame_listener.lookupTransform(self.world_frame, self.cmd_ref_frame, rospy.Time(0))

        self.world_trans = world_trans
        self.world_rot = world_rot
        # convert to base_link
        ref_to_world_tf = self.frame_listener.fromTranslationRotation(self.world_trans, self.world_rot );
        target_in_ref =  self.frame_listener.fromTranslationRotation(target_trans, target_rot);
        print 'ref_to_world_tf:\n ',ref_to_world_tf
        print 'target_in_ref:\n ', target_in_ref

        # shift the target from tool0 to link_6 since planning using tool0 do not works
        # just a shift in local reference frame, so use identity quaternion
        shift_in_ref = self.frame_listener.fromTranslationRotation(self.tool0_to_link6_trans, self.tool0_to_link6_rot)

#        print 'shift_in_table: \n',shift_in_table
        target_in_ref= numpy.dot(target_in_ref, shift_in_ref)
        print 'target_in_ref renewed:\n',target_in_ref
#        rospy.loginfo("After shifting to link_6: \n%s", target_in_table.__str__())

        target_in_world = numpy.dot(ref_to_world_tf, target_in_ref)
        print 'target_in_base: \n', target_in_world

        # retrieve target rotation and translation in base_link frame
        world_rot = tf.transformations.quaternion_from_matrix(target_in_world)
        world_trans = tf.transformations.translation_from_matrix(target_in_world)

        print '============================================='
        print 'target_in_world_rot:\n', world_rot
        print 'target_in_world_trans:\n', world_trans
        print 'target_in_world_rot_euler',tf.transformations.euler_from_quaternion(world_rot)

        pose_target = geometry_msgs.msg.Pose()
        pose_target.orientation = geometry_msgs.msg.Quaternion(*world_rot)
        pose_target.position = geometry_msgs.msg.Point(*world_trans)

#        rospy.loginfo("Target Pose:\n %s", pose_target.__str__())

#        rospy.loginfo("======= pickPointMsgToTargetPose END ======")
        return pose_target


# WARNING: all planning related function call all have to make sure the
# start state is the current state
    def onMoveArmService(self, req):

        pose_target = self.pickPointMsgToTargetPose(req.pick_point)
        try:
            robot_current_state = self.robot.get_current_state()
            #rospy.logerr("Current_state: %s",robot_current_state.__str__())
            self.group.set_start_state(robot_current_state)
            self.group.set_pose_target(pose_target)
            plan = self.group.plan()

            MAX_TRYS = 100
            count = 0
            # if the plan result is not valid ( cost too much time or
            # is empty, replan it
            while not self.is_plan_checked(plan):
                plan = self.group.plan()
                count += 1
                if count == MAX_TRYS:
                    raise Exception("No plan found after trying dead hard.")

#            rospy.loginfo(plan.__str__())

            rospy.loginfo("* onMoveInTableverticalService called")

            display_trajectory = moveit_msgs.msg.DisplayTrajectory()
            display_trajectory.trajectory_start= self.robot.get_current_state()
            display_trajectory.trajectory.append(plan)

            self.display_trajectory_publisher.publish(display_trajectory)

            # accelerate the trajectory by a factor of 2
            traj = plan.joint_trajectory
            for point in traj.points:
                point.time_from_start = rospy.Duration.from_sec(point.time_from_start.to_sec()/2.5)

#                rospy.loginfo("Excute Plan and Block... ")
            self.group.execute(plan)
            rospy.logerr("Excution Done! ")
            return True

        except Exception as e:
            print e
            return False
            pass

        pass

    def is_plan_checked(self,plan):

        # the maximum time allowed for execution, in seconds
        ALLOWED_TIME = 10.0
        ALLOWED_VARIANCE = 0.3

        if not plan:
            rospy.logerr("### Empty Plan Checked ")
            return False
        elif type(plan) is not moveit_msgs.msg.RobotTrajectory:
            rospy.logerr("### Plan Type invalid ")
            return False
        elif plan.joint_trajectory.points[len(plan.joint_trajectory.points)-1].time_from_start.to_sec() > ALLOWED_TIME:
            rospy.logerr("### Plan Time too long, giving up.")
            return False
        elif self.compute_variance(plan) > ALLOWED_VARIANCE:
            rospy.logerr("### Unexpected move may happen,give up.")
        else:
            rospy.logerr("### Planned Moving Time: %s", plan.joint_trajectory.points[len(plan.joint_trajectory.points)-1].time_from_start.to_sec())
            return True

    def compute_variance(self,plan):

        total = plan.joint_trajectory.points[len(plan.joint_trajectory.points)-1].time_from_start.to_sec()
        mean = total/(len(plan.joint_trajectory.points))
        durations = []
        total_difference = 0

        for index in range(len(plan.joint_trajectory.points)-2):
            duration = plan.joint_trajectory.points[index+1].time_from_start.to_sec()-plan.joint_trajectory.points[index].time_from_start.to_sec()
            difference = (duration - mean) * (duration - mean)
            total_difference += difference
        variance = total_difference/(len(plan.joint_trajectory.points))


        return variance


if __name__ == "__main__":
        # init rospy
        print "============== Starting Setup"
        moveit_commander.roscpp_initialize(sys.argv)
        rospy.init_node('move_in_table_node')

        # wait for moveit stuff starting up
        #rospy.sleep(10)

        node = MoveInTableNode()

        node.start();

        try:
            while not rospy.is_shutdown():
                rospy.sleep(1)
        except (KeyboardInterrupt,rospy.ROSInterruptException):
            pass

        pass
